/**
** @file:  operator_mpi.h
** @brief:
**/

#ifndef LATTICE_OPERATOR_MPI_H
#define LATTICE_OPERATOR_MPI_H

#include <mpi.h>
#include <complex>
#include <iostream>
#include "utils.h"

template <typename T>
double norm_2(const T &s)
{
    //    int rank;
    //    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    //if (rank==0){
    std::complex<double> s1(0.0, 0.0);
    for (int i = 0; i < s.size; i++) {
        s1 += s.A[i] * conj(s.A[i]);    // 这里利用虚数部分乘出来是 0 没有提升
    }
    double sum_n = s1.real();
    double sum;
    MPI_Reduce(&sum_n, &sum, 1, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
    MPI_Bcast(&sum, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Barrier(MPI_COMM_WORLD);

    return sum;

    //    return s1.real();
    //    }else
    //    {return 0;}
}

template <typename T>
double norm_2_half(const T &s)
{
    //    int rank;
    //    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    //if (rank==0){
    std::complex<double> s1(0.0, 0.0);
    for (int i = s.size/2; i < s.size; i++) {
        s1 += s.A[i] * conj(s.A[i]);    // 这里利用虚数部分乘出来是 0 没有提升
    }
    double sum_n = s1.real();
    double sum;
    MPI_Reduce(&sum_n, &sum, 1, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
    MPI_Bcast(&sum, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Barrier(MPI_COMM_WORLD);

    return sum;

    //    return s1.real();
    //    }else
    //    {return 0;}
}

// template <typename T>
// double norm_2_d(const T &s)
// {
//     double dsum = 0.0;
//     double* s_d = (double*)(s.A + i);
//     for (int i = 0; i < s.size; i++) {
//         dsum += s_d[0] * s_d[0];
//         dsum += s_d[1] * s_d[1];
//     }
//     double sum_n = dsum;
//     double sum;
//     MPI_Reduce(&sum_n, &sum, 1, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
//     MPI_Bcast(&sum, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);
//     MPI_Barrier(MPI_COMM_WORLD);

//     return sum;
// }

template <typename T>
std::complex<double> vector_p(const T &r1, const T &r2)
{
    // auto startTime = Clock::now();
    std::complex<double> s1(0.0, 0.0);
    // cblas_zdotc_sub(r1.size, r1.A, 1, r2.A, 1, &s1);
    for (int i = 0; i < r1.size; i++) {
        s1 += (conj(r1.A[i]) * r2.A[i]);    // 这里利用虚数部分乘出来是 0 没有提升
    }
    double sum_r = s1.real(); //fix
    double sum_i = s1.imag();
    double sumr;
    double sumi;
    // auto endTime = Clock::now();
    // auto compTime0 = std::chrono::duration_cast<std::chrono::microseconds>(endTime - startTime);
    // startTime = Clock::now();
    MPI_Reduce(&sum_r, &sumr, 1, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
    MPI_Reduce(&sum_i, &sumi, 1, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
    MPI_Bcast(&sumr, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Bcast(&sumi, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Barrier(MPI_COMM_WORLD);
    std::complex<double> sum(sumr, sumi);
    //sum.real()= sumr;
    //sum.imag()= sumi;
    // endTime = Clock::now();
    // auto compTime1 = std::chrono::duration_cast<std::chrono::microseconds>(endTime - startTime);
    // if(global_rank == 0)
    //     std::cout << "vector_p time: " << compTime0.count() << " " << compTime1.count() << std::endl;
    return sum;
};

// 对比上面，只能说速度几乎没有差别
template <typename T>
std::complex<double> vector_p_new(const T &r1, const T &r2)
{
    // auto startTime = Clock::now();
    // alignas(32) double sum4[4] = {0, 0, 0, 0};
    double ssum = 0;
    double* r1_double = (double*)(r1.A);
    double* r2_double = (double*)(r2.A);
    const int len = r1.size*2;
    for (int i = len/2; i < len; i++) {
        ssum += r1_double[i] * r2_double[i];
    }
    double sum_r = ssum;
    // double sum_r = sum4[0] + sum4[1] + sum4[2] + sum4[3]; //fix
    double sum_i = 0;
    double sumr;
    double sumi;
    // auto endTime = Clock::now();
    // auto compTime0 = std::chrono::duration_cast<std::chrono::microseconds>(endTime - startTime);
    // startTime = Clock::now();
    MPI_Reduce(&sum_r, &sumr, 1, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
    MPI_Reduce(&sum_i, &sumi, 1, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
    MPI_Bcast(&sumr, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Bcast(&sumi, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Barrier(MPI_COMM_WORLD);
    std::complex<double> sum(sumr, sumi);
    //sum.real()= sumr;
    //sum.imag()= sumi;
    // endTime = Clock::now();
    // auto compTime1 = std::chrono::duration_cast<std::chrono::microseconds>(endTime - startTime);
    // if(global_rank == 0)
    //     std::cout << "vector_p_new time: " << compTime0.count() << " " << compTime1.count() << std::endl;
    return sum;
};

template <typename T>
std::complex<double> vector_p_half(const T &r1, const T &r2)
{
    // auto startTime = Clock::now();
    std::complex<double> s1(0.0, 0.0);
    // cblas_zdotc_sub(r1.size, r1.A, 1, r2.A, 1, &s1);
    for (int i = r1.size/2; i < r1.size; i++) {
        s1 += (conj(r1.A[i]) * r2.A[i]);    // 这里利用虚数部分乘出来是 0 没有提升
    }
    double sum_r = s1.real(); //fix
    double sum_i = s1.imag();
    double sumr;
    double sumi;
    // auto endTime = Clock::now();
    // auto compTime0 = std::chrono::duration_cast<std::chrono::microseconds>(endTime - startTime);
    // startTime = Clock::now();
    MPI_Reduce(&sum_r, &sumr, 1, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
    MPI_Reduce(&sum_i, &sumi, 1, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
    MPI_Bcast(&sumr, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Bcast(&sumi, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Barrier(MPI_COMM_WORLD);
    std::complex<double> sum(sumr, sumi);
    //sum.real()= sumr;
    //sum.imag()= sumi;
    // endTime = Clock::now();
    // auto compTime1 = std::chrono::duration_cast<std::chrono::microseconds>(endTime - startTime);
    // if(global_rank == 0)
    //     std::cout << "vector_p time: " << compTime0.count() << " " << compTime1.count() << std::endl;
    return sum;
};

#endif //LATTICE_OPERATOR_MPI_H
